import torch
import torch.nn as nn
from modeling import deeplab
import cv2
from torchvision.transforms import transforms
import numpy as np
import math
from matplotlib import pyplot as plt
import imutils


class FixScaleCrop():
    def __init__(self, crop_size=513):
        super(FixScaleCrop, self).__init__()
        self.crop_size = crop_size

    def __call__(self, img):
        h, w = img.shape[:2]
        ratio = self.crop_size / min(h, w)
        n_h, n_w = int(h * ratio), int(w * ratio)
        img = cv2.resize(img, (n_w, n_h), interpolation=cv2.INTER_LINEAR)
        d_w = (n_w - self.crop_size) // 2
        d_h = (n_h - self.crop_size) // 2
        img = img[d_h:d_h + n_h, d_w:d_w + self.crop_size, :]
        return img


model = deeplab.DeepLab(num_classes=4)
model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load('deeplabv3.pkl'))
model.eval()
composed = transforms.Compose([
    FixScaleCrop(),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])


def deeplab_detect(origin: np.ndarray, label):
    global model
    h, w = origin.shape[:2]

    frame = composed(origin)
    frame = frame.unsqueeze(0)

    model = model.cuda()
    frame = frame.cuda()
    out = model(frame)
    out = out.squeeze_(0)
    out = torch.argmax(out, dim=0)
    out = out.cpu().numpy().astype(np.uint8)
    ratio = min(h, w) / 513

    n_h, n_w = int(ratio * 513), int(ratio * 513)
    out = cv2.resize(out, (n_w, n_h))
    dw, dh = (w - n_h) / 2, (h - n_h) / 2
    left = round(dw - 0.1)
    top = round(dh - 0.1)
    right = round(dw + 0.1)
    bottom = round(dh + 0.1)
    mask = cv2.copyMakeBorder(out, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
    kernel = np.ones(shape=(5, 5), dtype=np.uint8)
    mask = cv2.dilate(mask, kernel, iterations=1)
    if label == 2:
        mask_rotate = imutils.rotate(mask, -20, (mask.shape[0] // 2, mask.shape[1] // 2))
    elif label == 1:
        mask_rotate = imutils.rotate(mask, -20, (mask.shape[0] // 2, mask.shape[1] // 2))
    else:
        mask_rotate = mask
    # 绘制原来的轮廓
    conts = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    conts = imutils.grab_contours(conts)
    cv2.drawContours(origin, conts, 0, (0, 255, 0), 2)
    # 旋转后的轮廓
    conts = cv2.findContours(mask_rotate, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    conts = imutils.grab_contours(conts)
    conts = sorted(conts, key=cv2.contourArea, reverse=True)[0]

    M = cv2.moments(conts)
    cx = int(M['m10'] / M['m00'])
    cy = int(M['m01'] / M['m00'])
    hx, hy = sorted(conts, key=lambda x: x[0, 1])[0][0]

    if label == 2 or label == 0:
        if hx < cx and hy < cy:
            cv2.putText(origin, str(label + 1) + ':Close', (hx, hy), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        else:
            cv2.putText(origin, str(label + 1) + ':Open', (hx, hy), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    if label == 1:
        x, y, w, h = cv2.boundingRect(conts)
        # cv2.rectangle(mask_rotate, (x, y), (x + w, y + h), (255, 0, 0), 2)
        # cv2.imshow('', mask_rotate)
        # cv2.waitKey(0)
        # print("---------------", w, h)
        if w < 1.5 * h:
            cv2.putText(origin, str(label + 1) + ':Close', (hx, hy), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        else:
            cv2.putText(origin, str(label + 1) + ':Open', (hx, hy), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    return origin
